import subprocess
import logging
import re
from datetime import datetime

class HdcError(Exception):
    """
    Error class for hdc
    """
    pass

class Hdc:
    def __init__(self, path: str = "hdc"):
        self.path = path

    def targets(self) -> list[dict]:
        """
        Get hdc targets
            10.80.81.82:9999                TCP     Offline         unknown...      hdc
            192.168.100.1:8888              TCP     Connected       localhost       hdc
            1234567812345678                USB     Connected       localhost       hdc
            8765432187654321                USB     Offline         localhost       hdc

        Returns:
            list: Targets of hdc
        """
        cmd = f"{self.path} list targets -v"
        output = self.__send_command(cmd, 5)
        target_str_list = output.split("\n")
        target_list = []
        for target_str in target_str_list:
            target_fields = [field for field in target_str.split("\t") if field != '']
            logging.debug(f"target_fields: {target_fields}\n")
            if len(target_fields) >= 5 and target_fields[4] == "hdc":
                 target_list.append(
                    {
                        "device_id": target_fields[0].strip(),
                        "connect_type": target_fields[1].strip(),
                        "connect_status": target_fields[2].strip(),
                        "connect_host": target_fields[3].strip(),
                    })
        return target_list

    def __send_command_without_result(self, command: str):
        try:
            subprocess.run(
                command,
                capture_output=True,
                text=True,
                timeout=None
            )
        except subprocess.CalledProcessError:
            error_message = (
                f"Command '{command}' called process with error."
            )
            logging.error(error_message)
            raise HdcError(error_message)
        except FileNotFoundError:
            logging.error(f"HDC executable not found at '{self.path}'")
            raise

    def __send_command(self, command: str, timeout: float = 10) -> str:
        """
        Send command to hdc

        Args:
            command (str): Command to send
            timeout (float, optional): Timeout for command. Defaults to 10 seconds.

        Returns:
            str: Output of command
        """
        try:
            process = subprocess.run(
                command,
                capture_output=True,
                text=True,
                timeout=timeout
            )
            stdout = process.stdout.strip()        
            if process.returncode != 0:
                error_message = f"Command '{command}' failed with return code {process.returncode}."
                if stdout:
                    error_message += f"\nStdout: {stdout}"
                stderr = process.stderr.strip()
                if stderr:
                    error_message += f"\nStderr: {stderr}"
                logging.error(error_message)
                raise HdcError(error_message)
            return stdout
        except subprocess.TimeoutExpired:
            error_message = (
                f"Command '{command}' timed out after {timeout} seconds."
            )
            logging.error(error_message)
            raise HdcError(error_message)
        except FileNotFoundError:
            logging.error(f"HDC executable not found at '{self.path}'")
            raise

    def __send_shell_command(self, device_id: str, command: str, timeout: float = 10) -> str:
        """
        Send command to hdc

        Args:
            device (str): Device target tag
            command (str): Command to send
            timeout (float, optional): Timeout for command. Defaults to 10 seconds.

        Returns:
            str: Output of command
        """
        if not command.startswith("\"") or not command.endswith("\""):
            command = f"\"{command.strip("\"")}\""
        cmd = f"{self.path} shell {command}" \
                if device_id == '' \
                else f"{self.path} -t {device_id} shell ${command}"
        logging.debug(f"device[{device_id}] send command: {cmd}")
        result = self.__send_command(cmd, timeout)
        logging.debug(f"device[{device_id}] result: {result[:100]}...")
        return result
    
    def shell_no_return(self, device_id: str, command: str):
        if not command.startswith("\"") or not command.endswith("\""):
            command = f"\"{command.strip("\"")}\""
        cmd = f"{self.path} shell {command}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} shell ${command}"
        logging.debug(f"device[{device_id}] send command: {cmd}")
        self.__send_command_without_result(cmd)

    def shell(self, device_id: str, command: str, timeout: float = 10) -> str:
        """
        Execute shell command on device

        Args:
            device_id (str): Device serial number
            command (str): Command to execute
            timeout (float, optional): Timeout for command execution. Defaults to 10.

        Returns:
            str: Output of command
        """
        return self.__send_shell_command(device_id, command, timeout)

    def reboot(self, device_id: str):
        """
        Reboot device

        Args:
            device_id (str): Device serial number
        """
        self.__send_shell_command(device_id, "reboot")

    def version(self) -> str:
        """
        Get hdc version

        Returns:
            str: Version of hdc
        """
        cmd = f"{self.path} -v"
        result = self.__send_command(cmd, 2)
        return result.strip()

    def file_send(self, device_id: str, local_file: str, remote_file: str) -> dict:
        """
        Send file to device

        Args:
            device_id (str): Device serial number
            local_file (str): Local file path
            remote_file (str): Remote file path

        Returns:
            dict: File send result

                Success: {"error": { "code": 0, "message": None }, "result": { "size": 0, "file_count": 1, "time": 0, "rate": 0}}
                
                Fail:    {"error": { "code": -1, "message": "Error message"}, "result": None}
        """
        cmd = f"{self.path} file send {local_file} {remote_file}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} file send {local_file} {remote_file}"
        result = self.__send_command(cmd)
        if "[Fail]" in result:
            return {"error": { "code": -1, "message": result }, "result": None}
        else:
            pattern = r"Size:(?P<size>\d+),\s*File count\s*=\s*(?P<file_count>\d+),\s*time:(?P<time>\w+)\s*rate:(?P<rate>[\d.]+[\w/]+)"
            match = re.search(pattern, result)
            if match:
                size = match.group("size")
                file_count = match.group("file_count")
                time = match.group("time")
                rate = match.group("rate")
                return {"error": { "code": 0, "message": None }, "result": {"size": size, "file_count": file_count, "time": time, "rate": rate}}
            else:
                return {"error": { "code": -2, "message": "File send: Failed to parse the result" }, "result": None}

    def file_recv(self, device_id: str, remote_file: str, local_file: str) -> dict:
        """
        Receive file from device

        Args:
            device_id (str): Device serial number
            remote_file (str): Remote file path
            local_file (str): Local file path

        Returns:
            dict: Receive file result

                Success: {"error": { "code": 0, "message": None }, "result": { "size": 0, "file_count": 1, "time": 0, "rate": 0}}
                
                Fail:    {"error": { "code": -1, "message": "Error message"}, "result": None}
        """
        cmd = f"{self.path} file recv {remote_file} {local_file}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} file recv {remote_file} {local_file}"
        result = self.__send_command(cmd)
        if "[Fail]" in result:
            return {"error": { "code": -1, "message": result }, "result": None}
        else:
            pattern = r"Size:(?P<size>\d+),\s*File count\s*=\s*(?P<file_count>\d+),\s*time:(?P<time>\w+)\s*rate:(?P<rate>[\d.]+[\w/]+)"
            match = re.search(pattern, result)
            if match:
                size = match.group("size")
                file_count = match.group("file_count")
                time = match.group("time")
                rate = match.group("rate")
                return {"error": { "code": 0, "message": None }, "result": {"size": size, "file_count": file_count, "time": time, "rate": rate}}
            else:
                return {"error": { "code": -2, "message": "File recv: Failed to parse the result" }, "result": None}

    def file_recv_hilog(self, device_id: str, local_folder: str) -> dict:
        """
        Receive hilog file from device

        Args:
            device_id (str): Device serial number
            local_folder (str): Local folder path

        Returns:
            dict: File receive result
        """
        self.file_recv(device_id, "/data/log/hilog", local_folder)

    def file_recv_faultlog(self, device_id: str, local_folder: str) -> dict:
        """
        Receive faultlog from device, include faultlogger, freeze, temp

        Args:
            device_id (str): Device serial number
            local_folder (str): Local folder path

        Returns:
            dict: File receive result
        """
        self.file_recv(device_id, "/data/log/faultlog", local_folder)
    
    def file_recv_faultlogger(self, device_id: str, local_folder: str) -> dict:
        """
        Receive faultlogger file from device
        
        Args:
            device_id (str): Device serial number
            local_folder (str): Local folder path

        Returns:
            dict: File receive  result
        """
        return self.file_recv(device_id, "/data/log/faultlog/faultlogger", local_folder)
    
    def file_recv_memory_leak(self, device_id: str, local_folder: str) -> dict:
        """
        Receive memory leak log

        Args:
            device_id (str): Device serial number
            local_folder (str): Local folder path
        
        Returns:
            dict: File receive result
        """
        return self.file_recv(device_id, "/data/log/reliability/resource_leak/memory_leak", local_folder)
    
    def file_recv_fd_leak(self, device_id: str, local_folder: str) -> dict:
        """
        Receive file from device

        Args:
            device_id (str): Device serial number
            remote_file (str): Remote file path
            local_file (str): Local file path, could be empty

        Returns:
            dict: File receive result
        """
        return self.file_recv(device_id, "/data/log/reliability/resource_leak/fd_leak", local_folder)
    
    def file_recv_thread_leak(self, device_id: str, local_folder: str) -> dict:
        """
        Receive thread leak file from device
        
        Args:
            device_id (str): Device serial number
            local_folder (str): Local folder path

        Returns:
            dict: File receive result
        """
        return self.file_recv(device_id, "/data/log/reliability/resource_leak/thread_leak", local_folder)

    def fport_tcp_add(self, device_id: str, from_port: int, to_port: int) -> bool:
        """
        forward port

        Args:
            device_id (str): device id
            from_port (str): from port
            to_port (str): to port

        Returns::
            bool: True if success, False otherwise
        """
        cmd = f"{self.path} fport tcp:{from_port} tcp:{to_port}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} fport tcp:{from_port} tcp:{to_port}"
        result = self.__send_command(cmd)
        return True if "Forwardport result:OK" in result else False
    
    def fport_tcp_del(self, device_id: str, from_port: int, to_port: int) -> bool:
        """
        Delete a forward port

        Args:
           device_id: device id
           from_port: from port number
           to_port: to port number

        Returns:
           True or False
        """
        cmd = f"{self.path} fport rm tcp:{from_port} tcp:{to_port}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} fport rm tcp:{from_port} tcp:{to_port}"
        result = self.__send_command(cmd)
        return True if "Remove forward ruler success" in result else False

    def fport_list(self, device_id: str) -> list:
        """
        list forward port

        Args:
            device_id (str): device id

        Returns:
            list: list of forward port
        """
        cmd = f"{self.path} fport ls" \
            if device_id == '' \
            else f"{self.path} -t {device_id} fport ls"
        result = self.__send_command(cmd)
        return result.split("\n")
    
    def dump_heap(self, device_id: str, pid: int, tid: int = -1):
        """
        Dump arkts heap

        Args:
            device_id (str): Device serial number
            pid (int): Process id
        """
        if (pid <= 1):
            raise HdcError("pid can not less than 1")
        cmd = f"hidumper --mem-jsheap {pid} -T {tid}" \
            if tid > 0 \
            else f"hidumper --mem-jsheap {pid}"
        return self.__send_shell_command(device_id, cmd)

    def gc(self, device_id: str, pid: int, tid: int = -1):
        """
        Run gc of ArkVM

        Args:
            device_id (str): Device serial number
            pid (int): Process id
            tid (int, optional): Thread id. Defaults to -1.     
        """
        if (pid <= 1):
            raise HdcError("pid can not less than 1")
        cmd = f"hidumper --mem-jsheap {pid} -T {tid} --gc" \
            if tid > 0 \
            else f"hidumper --mem-jsheap {pid} --gc"
        return self.__send_shell_command(device_id, cmd)

    def mem_info(self, device_id: str, pid: int, tid: int = -1) -> dict:
        """
        Get memory info

        Args:
            device_id (str): Device ID
            pid (int): Process ID
            tid (int, optional): Thread ID. Defaults to -1.

        Raises:
            HdcError: If pid is less than 1

        Returns:
            dict: Memory info
            {
                "TimeStamp": int,
                "Total": "123456",
                "ark ts heap": "123456",
                "native heap": "123456",
                "AnonPage other": "123456",
                "FilePage other": "123456",
                "Graph": "123456",
                "GL": "123456",
                "jemalloc heap": "123456"
            }
        """
        if (pid <= 1):
            raise HdcError("pid can not less than 1")
        cmd = f"hidumper --mem {pid} -T {tid}" \
            if tid > 0 \
            else f"hidumper --mem {pid}"
        stdout = self.__send_shell_command(device_id, cmd)
        keys = ["Total", "ark ts heap", "native heap", "AnonPage other", "FilePage other", "Graph", "GL", "jemalloc heap:"]
        result = dict()
        result["TimeStamp"] = datetime.now().timestamp()
        for line in stdout.splitlines():
            for key in keys:
                loc = line.find(key)
                if loc == -1:
                    continue
                line = line[loc + len(key) :].strip()
                if line == '':
                    continue
                loc = line.find(" ")
                value = line[:loc]
                try:
                    intval = int(value)
                    result[key.strip().strip(":")] = intval
                except:
                    continue
        return result
        
    def api_version(self, device_id: str) -> str:
        """
        Get api version

        Args:
            device_id (str): Device serial number

        Returns:
            str: Api version
        """
        return self.__send_shell_command(device_id, f"param get const.ohos.apiversion")
    
    def software_version(self, device_id: str) -> str:
        """
        Get software version

        Args:
            device_id (str): Device serial number

        Returns:
            str: Software version
        """
        return self.__send_shell_command(device_id, f"param get const.product.software.version")
    
    def product_info(self, device_id: str) -> dict:
        """
        Get product info

        Args:
            device_id (str): Device serial number

        Returns:
            dict: 

            Example: {
                "product_name": "HUAWEI P30 Pro",
                "product_brand": "HUAWEI",
                "product_model": "HUAWEI P30 Pro",
                "product_version": "10",
                "product_api_version": "29",
                "product_cpu_abilist": cpu_abilist
            }          
        """
        name = self.__send_shell_command(device_id, f"param get const.product.name")
        brand = self.__send_shell_command(device_id, f"param get const.product.brand")
        model = self.__send_shell_command(device_id, f"param get const.product.model")
        software_version = self.__send_shell_command(device_id, f"param get const.product.software.version")
        cpu_abilist = self.__send_shell_command(device_id, f"param get const.product.cpu.abilist")
        api_version = self.__send_shell_command(device_id, f"param get const.ohos.apiversion")
        return {
            "product_name": name,
            "product_brand": brand,
            "product_model": model,
            "product_version": software_version,
            "product_api_version": api_version,
            "product_cpu_abilist": cpu_abilist
        }
    
    def pidof(self, device_id: str, bundle_name: str) -> int:
        """
        Get process id of bundle name

        Args:
            device_id (str): Device serial number
            bundle_name (str): Bundle name

        """
        result = self.__send_shell_command(device_id, f"pidof {bundle_name}")
        try:
            return int(result)
        except:  # noqa: E722
            return -1
        
    def get_pid(self, device_id: str, cmd_content: str):
        """
        Get process id of cmd_content

        Args:
            device_id (str): Device serial number
            cmd_content (str): Command content
        """
        result = self.__send_shell_command(device_id, f"ps -ef | grep '{cmd_content}' | grep -v grep")
        # UID     PID    PPID  C STIME     TTY      TIME     CMD
        # shell   10022  60065 0 13:28:14  ?        00:00:07 uitest uiRecord record
        patter = r"(?P<UID>\S+)\s+(?P<PID>\d+)\s+(?P<PPID>\d+)\s+(?P<C>\S+)\s+(?P<STIME>\S+)\s+(?P<TTY>\S+)\s+(?P<TIME>\S+)\s+(?P<CMD>[^\r\n\f]+)"
        match = re.match(patter, result)
        if match:
            return int(match.group("PID"))
        else:
            return -1

    def kill_pid(self, device_id: str, pid: int) -> bool:
        if (pid <= 0):
            raise HdcError(f"Invalid pid: {device_id} - {pid}")
        self.__send_shell_command(device_id, f"kill -9 {pid}")

    def installed_apps(self, device_id: str) -> list[str]:
        """
        Get installed apps

        Args:
            device_id (str): Device serial number

        """
        apps = self.__send_shell_command(device_id, "bm dump -a").split("\n")
        return [app.strip() for app in apps if ':' not in app]

    def install(self, device_id: str, hap_path: str):
        """
        Install hap package

        Args:
            device_id (str): Device serial number
            hap_path (str): Hap package path
        """
        cmd = f"{self.path} install {hap_path}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} install {hap_path}"
        self.__send_command(cmd)
    
    def uninstall(self, device_id: str, bundle_name: str):
        """
        Uninstall Hap package

        Args:
            device_id (str): Device serial number
            bundle_name (str): Hap bundle name
        """
        cmd = f"{self.path} uninstall {bundle_name}" \
            if device_id == '' \
            else f"{self.path} -t {device_id} uninstall {bundle_name}"
        self.__send_command(cmd)

    def aa_start(self, device_id: str, bundle_name: str, ability_name: str = "EntryAbility"):
        """
        Start app

        Args:
            device_id (str): Device serial number
            bundle_name (str): App bundle name
        """
        cmd = f"aa start -a {ability_name} -b {bundle_name}"
        self.__send_shell_command(device_id, cmd)

    def aa_stop(self, device_id: str, bundle_name: str):
        """
        Stop app

        Args:
            device_id (str): Device serial number
            bundle_name (str): App bundle name
        """
        cmd = f"aa force-stop {bundle_name}"
        self.__send_shell_command(device_id, cmd)

    def temperature(self, device_id: str) -> float | None:
        """
        Get device temperature

        Args:
            device_id (str): Device serial number

        Returns:
            float: Device temperature
        """
        cmd = f"hidumper -s BatteryService -a -i | grep temperature"
        result = self.__send_shell_command(device_id, cmd)
        pattern = r"temperature\s*:\s*(?P<temperature>\d+)"
        match = re.search(pattern, result)
        if match:
            return float(match.group("temperature")) / 10.0
        else:
            return None

    def cpu_count(self, device_id: str) -> int:
        """
        Get cpu count

        Args:
            device_id (str): Device serial number

        Returns:
            int: Cpu count, -1 if error
        """
        cmd = "cat /proc/cpuinfo | grep processor | wc -l"
        result = self.__send_shell_command(device_id, cmd)
        try :
            return int(result.strip())
        except ValueError :
            return -1


    def is_cpu_online(self, device_id: str, cpu_id: int) -> bool:
        """
        Only support device with root image

        Check whether the cpu is online

        Args:
            device_id (str): Device serial number
            cpu_id (int): CPU ID

        Returns:
            bool: True if the cpu is online, False otherwise
        """
        cmd = f"cat /sys/devices/system/cpu/cpu{cpu_id}/online"
        result = self.hdc.shell(device_id, cmd).strip()
        return result == "1"

    def set_cpu_online(self, device_id: str, cpu_id: int):
        """
        Only support device with root image

        Set cpu online

        Args:
            device_id (str): Device serial number
            cpu_id (int): CPU id
        """
        cmd = f"echo 1 > sys/devices/system/cpu/cpu{cpu_id}/online"
        self.__send_shell_command(device_id, cmd)

    def cpu_available_frequencies(self, device_id: str, cpu_id: int) -> list[int]:
        """
        Only support device with root image

        Get cpu available frequencies

        Args:
            device_id (str): Device serial number
            cpu_id (int): CPU id

        Returns:
            list[int]: CPU frequencies
        """
        cmd = f"cat sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_available_frequencies"
        result = self.__send_shell_command(device_id, cmd)
        frequencies = [int(f) for f in result.split()]
        return frequencies

    def get_cpu_frequency(self, device_id: str, cpu_id: int) -> dict[str, int]:
        """
        Only support device with root image

        Get CPU frequency

        Args:
            device_id (str): Device id
            cpu_id (int): CPU id

        Returns:
            dict: CPU frequency

            {
                "cpuinfo_cur_freq": int,
                "cpuinfo_max_freq": int,
                "cpuinfo_min_freq": int
            }
        """
        items = [
            "cpuinfo_cur_freq",
            "cpuinfo_max_freq",
            "cpuinfo_min_freq"
        ]
        result: dict = dict(str, int)
        for item in items:
            output = self.__send_shell_command(
                device_id,
                f"cat /sys/devices/system/cpu/cpu{cpu_id}/cpufreq/{item}"
            )
            try:
                result[item] = int(output.strip())
            except ValueError:
                raise HdcError(f"Failed to parse frequency for {item}: {output}")
        return result

    def get_all_cpu_freqs(self, device_id: str) -> list[dict[str, int]]:
        """
        get all cpu frequencies

        Args:
            device_id (str): device id

        Returns:
            list[dict[str, int]]: cpu frequencies

            Example:
            [
                {
                    "cpu_id": int,
                    "cpuinfo_cur_freq": int,
                    "cpuinfo_max_freq": int
                },
                ...
            ]
        """
        cpu_count = self.cpu_count(device_id)
        if cpu_count == -1:
            return {}
        
        result = []
        output = self.__send_shell_command(device_id, f"hidumper --cpufreq | grep -E '[0-9]+'")
        pattern = r"cmd is: cat /sys/devices/system/cpu/cpu(?P<cpu_id>\d+)/cpufreq/(?P<cpu_freq_type>\w+)\s*[\r\n]+\s*(?P<cpu_freq>\d+)"
        matches = re.findall(pattern, output)
        if len(matches) != cpu_count * 2:
            return {}
        else:
            for match in matches: # tuple ('0', 'cpuinfo_max_freq', '2600000')
                cpu_id = int(match[0].strip())
                cpu_freq_type = match[1].strip()
                cpu_freq = int(match[2].strip())       
                if (cpu_id < len(result)):
                    result[cpu_id][cpu_freq_type] = cpu_freq
                else:
                    result.append({ "cpu_id": cpu_id, cpu_freq_type: cpu_freq })
        return result

    def set_cpu_frequency(self, device_id: str, cpu_id: int, max_frequency: int, min_frequency: int):
        """
        Only support device with root image

        Set CPU frequency

        Args:
            device_id (str): Device ID
            cpu_id (int): CPU ID
            max_frequency (int): Max frequency, if value bigger than 0, it will be set to the maximum frequency.
            min_frequency (int): Min frequency, if value bigger than 0, it will be set to the minimum frequency.
        """
        if (max_frequency > 0):
            cmd = f"echo {max_frequency} > /sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_max_freq"
            self.__send_shell_command(device_id, cmd)
        if (min_frequency > 0):
            cmd = f"echo {min_frequency} > /sys/devices/system/cpu/cpu{cpu_id}/cpufreq/scaling_min_freq"
            self.__send_shell_command(device_id, cmd)


